import cv2
import numpy as np
import matplotlib.pyplot as plt

'''
    Brute-Force 蛮力匹配
'''
img1 = cv2.imread('box.png', 0)
img2 = cv2.imread('box_in_scene.png', 0)


def cv_show(name, img):
    cv2.imshow(name, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


cv_show('img1', img1)
cv_show('img2', img2)


sift = cv2.xfeatures2d.SIFT_create()
kp1, des1 = sift.detectAndCompute(img1, None)
kp2, des2 = sift.detectAndCompute(img2, None)
# crossCheck 表示两个特征点要相互匹配, 例如 A 中第 i 个特征点与 B 中第 j 个特征点最近的, 并且 B 中第 j 个特征点到 A 中第 i 个特征点也是最近的
# NORM_L2: 归一化数组的(鸥几里德距离), 如果其它特征计算方法需要考虑不同的匹配计算方式
bf = cv2.BFMatcher(crossCheck=True)

# 一对一的匹配
matches = bf.match(des1, des2)
# matches.sorted(matches, key=lambda x: x.distance)
matches.sort(key=lambda x: x.distance)
# 其中，img1 和 img2 是要绘制匹配点的两幅图像，kp1 和 kp2 是两幅图像的关键点，matches[:10] 是匹配点的列表，None 表示不需要输出的图像，flags=2 表示绘制匹配线。
img3 = cv2.drawMatches(img1, kp1, img2, kp2, matches[:10], None, flags=2)
cv_show('img3', img3)

# k对最佳匹配
bf = cv2.BFMatcher()
matches = bf.knnMatch(des1, des2, k=2)
good = []
for m, n in matches:
    if m.distance < 0.75 * n.distance:
        good.append([m])
img3 = cv2.drawMatchesKnn(img1, kp1, img2, kp2, matches[:10], None, flags=2)
cv_show('img3', img3)

